import streamlit as st
from dotenv import load_dotenv
import os
from database import get_messages, add_message, update_conversation_timestamp, add_conversation
from auth import get_current_user, is_logged_in
from utils.util import image_to_base64
from openai import OpenAI
# Agent 相关开始
import pprint
from qwen_agent.agents import Assistant
from agent.agent_tools import MyImageGen, GetCurrentDatetime, GetOneMemberIntroduction
import json

# 加载 .env 文件
load_dotenv()

# 获取 API URL
base_url = os.getenv("virtaicloud_base_url")
api_key = os.getenv("virtaicloud_api_key")
modelname = os.getenv("virtaicloud_modelname")
    
# 步骤 2：配置您所使用的 LLM。
llm_cfg = {
    # 使用 DashScope 提供的模型服务：
    'model': modelname,
    'model_server': base_url,
    'api_key': api_key,
    # 'api_key': 'YOUR_DASHSCOPE_API_KEY',
    # 如果这里没有设置 'api_key'，它将读取 `DASHSCOPE_API_KEY` 环境变量。

    # 使用与 OpenAI API 兼容的模型服务，例如 vLLM 或 Ollama：
    # 'model': 'Qwen2-7B-Chat',
    # 'model_server': 'http://localhost:8000/v1',  # base_url，也称为 api_base
    # 'api_key': 'EMPTY',

    # （可选） LLM 的超参数：
    'generate_cfg': {
        'top_p': 0.8
    }
}

# 步骤 3：创建一个智能体。这里我们以 `Assistant` 智能体为例，它能够使用工具并读取文件。
system_instruction = '''你是一个乐于助人的AI助手。
在收到用户的请求后，你应该：
你总是用中文回复用户。'''
tools = ['my_image_gen', 'get_current_datetime','get_one_member_introduction', 'code_interpreter']  # `code_interpreter` 是框架自带的工具，用于执行代码。
files = ['./examples/resource/doc.pdf']  # 给智能体一个 PDF 文件阅读。
bot = Assistant(llm=llm_cfg,
                system_message=system_instruction,
                function_list=tools,
                # files=files
               )

# Agent 相关结束

# 获取当前路径
curr_path = os.getcwd()
# 创建 upload 文件夹（如果不存在）
upload_dir = "upload"
os.makedirs(upload_dir, exist_ok=True)

# st.write("base_url", base_url)
# st.write("api_key", api_key)
# st.write("modelname", modelname)

# 趋动云
client = OpenAI(
    api_key=api_key,
    base_url=base_url
)

# 请求 Qwen 大模型
def get_qwen_response(user_input, image_path=None):
    try:
        messages = [{"role": "user", "content": []}]

        # 如果有图片文件，添加到消息中
        if image_path:
            # st.write("image_path", os.path.join(curr_path,image_path))
            base64_image = image_to_base64(os.path.join(curr_path,image_path))
            image_url = f"data:image/png;base64,{base64_image}"
            messages[0]["content"].append({"type": "image_url", "image_url": {"url": image_url}})

        # 添加文本消息
        if user_input:
            messages[0]["content"].append({"type": "text", "text": user_input})

        # st.write(messages)
        # st.write(modelname)
        # 请求 大模型
        completion = client.chat.completions.create(
            model=modelname,
            messages=messages
        )
        return completion.choices[0].message.content
    except Exception as e:
        st.error(f"请求大模型失败: {e}")
        return "请求失败，请稍后重试。"

messages = []  # 这里储存聊天历史。
# 修改 get_model_response 函数以使用 Qwen 大模型
def get_model_response(user_input, image_path=None):
    # 直接请求大模型
    # return get_qwen_response(user_input, image_path)

    # 调用 Agent
    # 将用户请求添加到聊天历史。
    messages.append({'role': 'user', 'content': user_input})
    response = bot.run_nonstream(messages=messages)
    # response = []
    # for response in bot.run(messages=messages):
        # 流式输出。
        # print('机器人回应:')
        # pprint.pprint(response, indent=2)
        # return response
        # response.append(response)
    # 将机器人的回应添加到聊天历史。
    print('机器人回应:', response)
    messages.extend(response)
    return response

# 自定义消息框样式
def message_box(message, is_user=False):
    if is_user:
        # 用户消息：右对齐，带背景色，右侧显示用户图标
        background_color = "#EFF6FF"  # RGB(239, 246, 255) 的十六进制值
        align = "right"
        st.markdown(
            f"""
            <div style="
                display: flex;
                justify-content: flex-end;
                align-items: center;
                margin: 5px 0;
            ">
                <div style="
                    background-color: {background_color};
                    padding: 10px;
                    border-radius: 10px;
                    max-width: 70%;
                    text-align: {align};
                ">
                    {message}
                </div>
                <div style="
                    width: 30px;
                    height: 30px;
                    margin-left: 10px;
                    border-radius: 50%;
                    background-color: #4CAF50;
                    color: white;
                    display: flex;
                    justify-content: center;
                    align-items: center;
                    font-weight: bold;
                ">
                    U
                </div>
            </div>
            """,
            unsafe_allow_html=True,
        )
    else:
        # 模型消息：左对齐，无背景色，左侧显示模型图标
        align = "left"
        col1, col2 = st.columns([0.11, 5])  # 使用列布局放置图标和消息
        with col1:
            st.image("assets/logo.png", width=30)  # 加载 logo.png
        with col2:
            st.markdown(
                f"""
                <div style="
                    padding: 10px;
                    border-radius: 10px;
                    max-width: 100%;
                    text-align: {align};
                ">
                    {message}
                </div>
                """,
                unsafe_allow_html=True,
            )

# 校验输入内容
def validate_input(user_input):
    # 检查输入是否为空或仅包含空格
    if not user_input or not user_input.strip():
        return False
    return True

# 保存上传的图片
def save_uploaded_file(uploaded_file):
    # 生成唯一的文件名
    import uuid
    file_name = os.path.join(upload_dir, f"{uuid.uuid4()}_{uploaded_file.name}")
    with open(file_name, "wb") as f:
        f.write(uploaded_file.getbuffer())
    return file_name

# 自定义消息框样式
def message_box(message, is_user=False):
    if is_user:
        background_color = "#EFF6FF"
        align = "right"
        st.markdown(
            f"""
            <div style="display: flex; justify-content: flex-end; align-items: center; margin: 5px 0;">
                <div style="background-color: {background_color}; padding: 10px; border-radius: 10px; max-width: 70%; text-align: {align};">
                    {message}
                </div>
                <div style="width: 30px; height: 30px; margin-left: 10px; border-radius: 50%; background-color: #4CAF50; color: white; display: flex; justify-content: center; align-items: center; font-weight: bold;">
                    U
                </div>
            </div>
            """,
            unsafe_allow_html=True,
        )
    else:
        align = "left"
        col1, col2 = st.columns([0.11, 5])
        with col1:
            st.image("assets/logo.png", width=30)
        with col2:
            st.markdown(
                f"""
                <div style="padding: 10px; border-radius: 10px; max-width: 100%; text-align: {align};">
                    {message}
                </div>
                """,
                unsafe_allow_html=True,
            )

# 解析消息内容并显示图片
def parse_and_display_message(content):
    # 检查消息内容是否包含 {image_url: ...} 格式
    if "{image_url:" in content:
        try:
            # 提取 {image_url: ...} 部分
            json_start = content.find("{")
            json_end = content.find("}") + 1
            json_str = content[json_start:json_end]
            
            # 将 {image_url: ...} 转换为合法的 JSON 格式
            json_str = json_str.replace("image_url", '"image_url"')  # 添加双引号
            data = json.loads(json_str)  # 解析 JSON
            image_url = data["image_url"]
            
            # 显示图片，并支持点击在新窗口中打开
            st.markdown(
                f'<a href="{image_url}" target="_blank">'
                f'<img src="{image_url}" width="300" style="border-radius: 10px; margin: 10px 0;">'
                f'</a>',
                unsafe_allow_html=True
            )
            return  # 如果是图片消息，直接返回，不显示文本内容
        except json.JSONDecodeError:
            pass  # 如果 JSON 解析失败，继续显示文本内容
    # 检查消息内容是否包含 {current_time: ...} 格式
    elif "{current_time:" in content:
        try:
            # 提取 } 后面的内容
            json_end = content.find("}") + 1
            display_content = content[json_end:].strip()  # 去掉前后空格
            st.markdown(display_content)  # 显示 } 后面的内容
            return  # 如果是时间消息，直接返回，不显示其他内容
        except Exception as e:
            st.error(f"解析消息失败: {e}")
            return
    # 显示普通文本内容
    st.markdown(content)

# 渲染右侧主界面
def render_main(selected_conversation):
    # 检查用户是否登录
    if is_logged_in():
        user = get_current_user()
        # 开始 right-column div
        st.markdown('<div class="right-column">', unsafe_allow_html=True)

        # 添加标题和提示文字
        st.markdown(
            """
            <h2 style='text-align: center;'>我是 Doraemon-Bot AI 助手，很高兴见到你！</h2>
            <p style='text-align: center;'>我可以协助你解决客户服务中遇到的问题，请把你遇到的难题交给我吧~</p>
            """,
            unsafe_allow_html=True,
        )

        # 如果是新会话模式
        if selected_conversation == "new_conversation":
            st.info("这是一个新会话，请输入你的问题。")

        # 创建两个容器：一个用于消息显示，一个用于输入框
        message_container = st.container()  # 消息显示区域
        response_placeholder = st.empty()  # 创建一个占位符用于显示流式输出
        
        input_container = st.container()    # 输入框区域
        # 在消息显示区域加载历史消息
        with message_container:
            if selected_conversation != "new_conversation":
                messages = get_messages(selected_conversation)
                for msg in messages:
                    role, content, image_path, created_at = msg
                    # 如果role是用户或者admin，消息右对齐，带背景色
                    if role in ["user", "admin"]:
                        message_box(f"{content}", is_user=True)  # 用户消息右对齐，带背景色
                        if image_path:  # 如果消息包含图片，显示图片
                            st.image(image_path, width=300)
                    else:
                        # message_box(f"{content}")  # 模型消息左对齐，无背景色
                        # 解析并显示消息内容
                        print(content)
                        parse_and_display_message(content)

        # 在输入框区域固定输入框
        with input_container:
            st.markdown(
                """
                <style>
                .fixed-bottom {
                    position: fixed;
                    bottom: 0;
                    width: 100%;
                    background-color: white;
                    padding: 10px;
                    box-shadow: 0 -2px 10px rgba(0, 0, 0, 0.1);
                    z-index: 1000;
                    display: flex; /* 使用flex布局 */
                    align-items: center; /* 垂直居中对齐 */
                }
                .upload-container, .send-btn {
                    margin: 0 5px; /* 给左右两边添加一点间距 */
                }
                /* 为带有特定key的按钮设置样式 */
                div[data-baseweb="button"][key="send_btn"] button {
                    width: 80px;
                    height: 80px;
                }
                </style>
                <div class="fixed-bottom">
                """,
                unsafe_allow_html=True,
            )

            # 输入框
            if "input_value" not in st.session_state:
                st.session_state.input_value = ""  # 初始化输入框的值

            user_input = st.text_area(
                "输入消息",
                height=100,
                key="input_box",
                value=st.session_state.input_value,  # 绑定输入框的值到 session_state
            )
        # 将上传图片和发送按钮放在同一行
        col1, col2 = st.columns([3, 3])  # 3: 上传按钮宽度，3: 发送按钮宽度
        with col1:
            # 上传图片组件
            uploaded_file = st.file_uploader(
                "上传图片",
                type=["png", "jpg", "jpeg"],
                key="file_uploader",
                label_visibility="collapsed",  # 隐藏标签
            )
        with col2:
            # 发送按钮
            if st.button("发送", key="send_btn", help="发送消息"):
                # 校验输入内容
                if not validate_input(user_input) and not uploaded_file:
                    st.error("输入内容不能为空！")
                else:
                    # 获取当前用户
                    user = get_current_user()

                    if selected_conversation == "new_conversation":
                        # 如果是新会话，创建新对话
                        new_conversation_name = f"对话 {len(get_messages(user['id'])) + 1}"
                        conversation_id = add_conversation(user["id"], new_conversation_name)
                        st.session_state.selected_conversation = conversation_id
                    else:
                        # 如果不是新会话，使用当前选中的对话
                        conversation_id = selected_conversation

                    # 保存用户消息
                    if validate_input(user_input):
                        add_message(conversation_id, user.get("role"), user_input)

                    # 如果有图片文件，保存图片
                    image_path = None
                    if uploaded_file:
                        # 保存图片到 upload 文件夹
                        image_path = save_uploaded_file(uploaded_file)
                        # 将图片路径保存到数据库
                        add_message(conversation_id, user.get("role"), "图片已上传", image_path=image_path)

                    # 获取模型回复并逐步显示
                    
                    full_response = ""  # 用于保存完整的模型响应

                    # 获取模型流式响应
                    for chunk in get_model_response(user_input, image_path):
                        # 如果 chunk 是字典，提取文本内容
                        if isinstance(chunk, dict):
                            chunk = chunk.get("content", "")  # 假设字典中有一个 "content" 字段
                        # 确保 chunk 是字符串
                        chunk = str(chunk)
                        full_response += chunk  # 将流式输出的每一部分拼接到完整响应中
                        response_placeholder.markdown(full_response)  # 更新占位符内容

                    # 保存完整的模型响应到数据库
                    add_message(conversation_id, "assistant", full_response)

                    # 更新会话的最后更新时间
                    update_conversation_timestamp(conversation_id)

                    # 清空输入框
                    st.session_state.input_value = ""

                    # 刷新页面以显示新消息
                    st.rerun()

        st.markdown("</div>", unsafe_allow_html=True)  # 关闭 fixed-bottom div

        # 在对话框下方添加提示文字
        st.markdown(
            """
            <p style='font-size: 12px; text-align: center; color: #666;'>
                内容由 AI 生成，请仔细甄别
            </p>
            """,
            unsafe_allow_html=True,
        )

        # 结束 right-column div
        st.markdown('</div>', unsafe_allow_html=True)
